import cv2
import threading
from flask import Flask, Response, request
from flask_socketio import SocketIO
from ultralytics import YOLO
import time
import os
import logging
from flask_cors import CORS

# 设置详细日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

app = Flask(__name__)
CORS(app)  # 允许所有跨域请求
socketio = SocketIO(
    app, 
    cors_allowed_origins="*", 
    async_mode='threading', 
    logger=True, 
    engineio_logger=True,
    ping_timeout=60,  # 增加超时时间
    ping_interval=25  # 增加ping间隔
)

# 加载模型
logging.info("正在加载模型...")
try:
    model_path = os.path.join(os.path.dirname(__file__), 'models', 'best1.onnx')
    model = YOLO(model_path, task='detect')
    logging.info(f"模型加载成功: {model_path}")
except Exception as e:
    logging.error(f"模型加载失败: {str(e)}")
    model = None

# 全局变量
camera = None
is_running = False
frame_counter = 0
start_time = time.time()
last_frame_time = time.time()

def generate_frames():
    # 修复：将所有全局变量声明放在函数开头
    global camera, is_running, frame_counter, start_time, last_frame_time
    logging.info("视频流生成器启动")
    
    # 添加重试机制
    retry_count = 0
    max_retries = 5
    
    while is_running and retry_count < max_retries:
        try:
            if camera is None or not camera.isOpened():
                logging.warning("摄像头未初始化或未打开! 尝试重新初始化...")
                try:
                    # 尝试重新初始化摄像头
                    camera = cv2.VideoCapture(0)
                    if not camera.isOpened():
                        for i in range(1, 5):
                            camera = cv2.VideoCapture(i)
                            if camera.isOpened():
                                break
                    
                    if camera.isOpened():
                        camera.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
                        camera.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
                        logging.info("摄像头重新初始化成功")
                    else:
                        logging.error("摄像头重新初始化失败")
                        retry_count += 1
                        time.sleep(1)
                        continue
                except Exception as e:
                    logging.error(f"摄像头重新初始化失败: {str(e)}")
                    retry_count += 1
                    time.sleep(1)
                    continue
                
            success, frame = camera.read()
            if not success:
                logging.warning("无法从摄像头读取帧! 尝试重新初始化...")
                retry_count += 1
                time.sleep(0.5)
                continue
            else:
                retry_count = 0  # 重置重试计数器
            
            current_time = time.time()
            elapsed_time = current_time - last_frame_time
            last_frame_time = current_time
            
            if frame_counter % 5 == 0:
                logging.debug(f"处理第 {frame_counter} 帧...")
                annotated_frame = frame.copy()
                
                if model:
                    try:
                        results = model(frame, imgsz=640, conf=0.5, device='cpu')
                        annotated_frame = results[0].plot()
                        
                        fire_count = sum(1 for box in results[0].boxes if results[0].names[int(box.cls)] == 'fire')
                        smoke_count = sum(1 for box in results[0].boxes if results[0].names[int(box.cls)] == 'smoke')
                        
                        total_elapsed = current_time - start_time
                        fps = frame_counter / total_elapsed if total_elapsed > 0 else 0
                        
                        socketio.emit('detection_stats', {
                            'fire': fire_count,
                            'smoke': smoke_count,
                            'fps': fps
                        })
                    except Exception as e:
                        logging.error(f"模型推理错误: {str(e)}")
            
            ret, buffer = cv2.imencode('.jpg', annotated_frame)
            if not ret:
                logging.warning("帧编码失败!")
                continue
                
            frame_bytes = buffer.tobytes()
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')
            
            frame_counter += 1
            
            # 控制帧率
            target_fps = 15
            sleep_time = max(0, (1/target_fps) - elapsed_time)
            time.sleep(sleep_time)
            
        except Exception as e:
            logging.error(f"视频流生成器异常: {str(e)}")
            retry_count += 1
            time.sleep(1)
            
    logging.warning(f"视频流生成器因连续失败退出，重试次数: {retry_count}")
    if retry_count >= max_retries:
        socketio.emit('error', {'message': '视频流生成失败，请检查摄像头连接'})
        is_running = False  # 修复：移除多余的global声明

@app.route('/video_feed')
def video_feed():
    logging.info("接收到视频流请求")
    if not is_running:
        return Response("请先启动检测", status=400)
    return Response(generate_frames(), 
                    mimetype='multipart/x-mixed-replace; boundary=frame')

@socketio.on('start_detection')
def handle_start():
    global camera, is_running, frame_counter, start_time, last_frame_time
    logging.info("接收到 start_detection 事件")
    
    if is_running:
        logging.warning("检测已在运行中")
        socketio.emit('status', {'status': 'already_running'})
        return
        
    logging.info("正在初始化摄像头...")
    try:
        # 尝试多个摄像头索引
        camera_indexes = [0, 1, 2, 3, 4]
        for i in camera_indexes:
            camera = cv2.VideoCapture(i)
            if camera.isOpened():
                logging.info(f"成功打开摄像头索引 {i}")
                break
            else:
                logging.warning(f"摄像头索引 {i} 打开失败")
        
        if not camera.isOpened():
            logging.error("无法打开摄像头！请检查摄像头连接和权限")
            socketio.emit('error', {'message': '无法访问摄像头，请检查连接和权限'})
            return
        
        logging.info("摄像头初始化成功")
        camera.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        camera.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        camera.set(cv2.CAP_PROP_FPS, 15)
        
        width = camera.get(cv2.CAP_PROP_FRAME_WIDTH)
        height = camera.get(cv2.CAP_PROP_FRAME_HEIGHT)
        fps = camera.get(cv2.CAP_PROP_FPS)
        logging.info(f"摄像头分辨率: {width}x{height}, FPS: {fps}")
        
        is_running = True
        frame_counter = 0
        start_time = time.time()
        last_frame_time = time.time()
        
        # 立即发送状态更新
        socketio.emit('status', {'status': 'detection_started'})
        logging.info("检测已启动")
        
    except Exception as e:
        logging.error(f"摄像头初始化失败: {str(e)}", exc_info=True)
        socketio.emit('error', {'message': f'摄像头初始化失败: {str(e)}'})

@socketio.on('stop_detection')
def handle_stop():
    global camera, is_running
    logging.info("接收到 stop_detection 事件")
    
    if is_running:
        is_running = False
        if camera is not None:
            try:
                camera.release()
                logging.info("摄像头已释放")
            except Exception as e:
                logging.error(f"摄像头释放失败: {str(e)}")
            camera = None
        socketio.emit('status', {'status': 'detection_stopped'})
        logging.info("检测已停止")

@app.route('/test_camera')
def test_camera():
    try:
        cap = cv2.VideoCapture(0)
        if not cap.isOpened():
            for i in range(1, 5):
                cap = cv2.VideoCapture(i)
                if cap.isOpened():
                    break
            
        if cap.isOpened():
            success, frame = cap.read()
            cap.release()
            if success:
                return "摄像头测试成功！摄像头已连接且可读取帧。"
            else:
                return "摄像头已连接但无法读取帧。"
        else:
            return "无法打开摄像头。请检查连接和权限。"
    except Exception as e:
        return f"摄像头测试失败: {str(e)}"

@app.route('/')
def index():
    return """
    <h1>火焰烟雾检测系统后端</h1>
    <p>可用端点:</p>
    <ul>
        <li><a href="/test_camera">/test_camera</a> - 测试摄像头</li>
        <li><a href="/video_feed">/video_feed</a> - 视频流 (需要先启动检测)</li>
    </ul>
    """

# 添加Socket.IO连接事件
@socketio.on('connect')
def handle_connect():
    logging.info(f"客户端已连接: {request.sid}")

@socketio.on('disconnect')
def handle_disconnect():
    logging.info(f"客户端已断开连接: {request.sid}")

if __name__ == '__main__':
    port = 5000
    logging.info(f"服务器启动: http://0.0.0.0:{port}")
    logging.info("测试摄像头是否可用: http://localhost:5000/test_camera")
    try:
        # 使用allow_unsafe_werkzeug=True解决开发环境下的安全问题
        socketio.run(
            app, 
            host='0.0.0.0', 
            port=port, 
            debug=False,
            allow_unsafe_werkzeug=True
        )
    except OSError as e:
        if "Address already in use" in str(e):
            port = 5001
            logging.warning(f"端口5000被占用，尝试使用端口{port}")
            socketio.run(
                app, 
                host='0.0.0.0', 
                port=port, 
                debug=False,
                allow_unsafe_werkzeug=True
            )